package com.vincent.llm;

import cn.hutool.core.io.resource.ResourceUtil;
import com.alibaba.fastjson.JSON;
import com.alibaba.fastjson.JSONObject;

public class MaxModelLenCalculator2 {
    public static long calculateMaxModelLen(
            double gpuMemoryGB,        // 单个GPU显存（GB）
            int numGPUs,               // GPU数量
            int tensorParallelSize,    // 张量并行度
            double gpuMemoryUtilization, // 显存利用率（0到1）
            String dtype,              // 数据类型（"fp16"或"fp32"）
            long modelParams,          // 模型参数量
            int hiddenSize,            // 隐藏层大小
            int numLayers              // 层数
    ) {
        // 根据数据类型确定每个参数的字节数
        int bytesPerDtype = "fp16".equalsIgnoreCase(dtype) ? 2 : 4;

        // 将显存转换为字节（1GB = 1024^3 bytes）
        long gpuMemoryBytes = (long) (gpuMemoryGB * 1024 * 1024 * 1024);

        // 每个GPU的可用显存（显存 * 利用率）
        long availablePerGPU = (long) (gpuMemoryBytes * gpuMemoryUtilization);

        // 计算每个GPU的权重显存占用
        long weightPerGPU = modelParams * bytesPerDtype / tensorParallelSize;

        // 如果权重显存超过可用显存，直接返回0
        if (weightPerGPU > availablePerGPU) {
            return 0;
        }

        // 计算剩余可用显存（扣除权重后）
        long remainingMemory = availablePerGPU - weightPerGPU;

        // 计算分子（剩余显存 * 张量并行度）
        long numerator = remainingMemory * tensorParallelSize;

        // 分母（隐藏层大小 * 层数 * 字节数）
        long denominator = (long) (hiddenSize * numLayers * bytesPerDtype);

        // 如果分母为0（隐藏层或层数为0），抛出异常
        if (denominator == 0) {
            throw new IllegalArgumentException("隐藏层大小或层数不能为0");
        }

        // 计算最大序列长度（向下取整）
        long maxModelLen = numerator / denominator;

        return maxModelLen;
    }

    public static void main(String[] args) {
        String cfgStr = ResourceUtil.readUtf8Str("config.json");
        JSONObject cfg = JSON.parseObject(cfgStr);
        // 示例参数：
        double gpuMemoryGB = 40.0;        // 单个GPU显存16GB
        int numGPUs = 2;                  // 使用4个GPU
        int tensorParallelSize = 2;       // 张量并行度4
        double gpuMemoryUtilization = 0.95; // 显存利用率80%
        int maxPositionEmbeddings = cfg.getIntValue("max_position_embeddings");
        System.out.println("默认最大模型长度:"+maxPositionEmbeddings);
        String dtype = "fp16";            // 使用FP16
        long modelParams = 7 * 1000_000_000L; // 模型参数量175B
        int hiddenSize = cfg.getIntValue("hidden_size");            // 隐藏层大小4096
        int numLayers = cfg.getIntValue("num_hidden_layers");               // 32层

        long maxLen = calculateMaxModelLen(
                gpuMemoryGB, numGPUs, tensorParallelSize,
                gpuMemoryUtilization, dtype, modelParams,
                hiddenSize, numLayers
        );
        maxLen  = (long) (maxLen * 0.9f);

        System.out.println("最大模型长度: " + Math.min(maxLen, maxPositionEmbeddings));
    }
}
